// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/types/type.h>

#include <algorithm>
#include <climits>
#include <cstddef>
#include <iterator>
#include <limits>
#include <memory>
#include <mutex>
#include <ostream>
#include <sstream>  // IWYU pragma: keep
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include <nebula/core/array.h>
#include <nebula/core/compare.h>
#include <nebula/core/record_batch.h>

#include <turbo/utility/status.h>
#include <nebula/core/table.h>
#include <turbo/base/checked_cast.h>
#include <nebula/numeric/decimal.h>
#include <nebula/util/hash_util.h>
#include <nebula/util/hashing.h>
#include <nebula/util/key_value_metadata.h>
#include <turbo/log/logging.h>
#include <nebula/util/range.h>
#include <nebula/util/string.h>

#include <nebula/util/vector.h>
#include <nebula/core/visit_type_inline.h>

namespace nebula {

    std::vector<Type::type> AllTypeIds() {
        return {Type::NA,
                Type::BOOL,
                Type::INT8,
                Type::INT16,
                Type::INT32,
                Type::INT64,
                Type::UINT8,
                Type::UINT16,
                Type::UINT32,
                Type::UINT64,
                Type::FP16,
                Type::FP32,
                Type::FP64,
                Type::DECIMAL128,
                Type::DECIMAL256,
                Type::DATE32,
                Type::DATE64,
                Type::TIME32,
                Type::TIME64,
                Type::TIMESTAMP,
                Type::INTERVAL_DAY_TIME,
                Type::INTERVAL_MONTHS,
                Type::DURATION,
                Type::STRING,
                Type::BINARY,
                Type::LARGE_STRING,
                Type::LARGE_BINARY,
                Type::STRING_VIEW,
                Type::BINARY_VIEW,
                Type::FIXED_SIZE_BINARY,
                Type::STRUCT,
                Type::LIST,
                Type::LARGE_LIST,
                Type::LIST_VIEW,
                Type::LARGE_LIST_VIEW,
                Type::FIXED_SIZE_LIST,
                Type::MAP,
                Type::DENSE_UNION,
                Type::SPARSE_UNION,
                Type::DICTIONARY,
                Type::EXTENSION,
                Type::INTERVAL_MONTH_DAY_NANO,
                Type::RUN_END_ENCODED};
    }

    namespace internal {

        struct TypeIdToTypeNameVisitor {
            std::string out;

            template<typename ArrowType>
            turbo::Status Visit(const ArrowType *) {
                out = ArrowType::type_name();
                return turbo::OkStatus();
            }
        };

        std::string to_type_name(Type::type id) {
            TypeIdToTypeNameVisitor visitor;

            KCHECK_OK(visit_type_id_inline(id, &visitor));
            return std::move(visitor.out);
        }

        std::string to_string(Type::type id) {
            switch (id) {
#define TO_STRING_CASE(_id) \
      case Type::_id:           \
        return TURBO_STRINGIFY(_id);

                TO_STRING_CASE(NA)
                TO_STRING_CASE(BOOL)
                TO_STRING_CASE(INT8)
                TO_STRING_CASE(INT16)
                TO_STRING_CASE(INT32)
                TO_STRING_CASE(INT64)
                TO_STRING_CASE(UINT8)
                TO_STRING_CASE(UINT16)
                TO_STRING_CASE(UINT32)
                TO_STRING_CASE(UINT64)
                TO_STRING_CASE(FP16)
                TO_STRING_CASE(FP32)
                TO_STRING_CASE(FP64)
                TO_STRING_CASE(DECIMAL128)
                TO_STRING_CASE(DECIMAL256)
                TO_STRING_CASE(DATE32)
                TO_STRING_CASE(DATE64)
                TO_STRING_CASE(TIME32)
                TO_STRING_CASE(TIME64)
                TO_STRING_CASE(TIMESTAMP)
                TO_STRING_CASE(INTERVAL_DAY_TIME)
                TO_STRING_CASE(INTERVAL_MONTH_DAY_NANO)
                TO_STRING_CASE(INTERVAL_MONTHS)
                TO_STRING_CASE(DURATION)
                TO_STRING_CASE(STRING)
                TO_STRING_CASE(STRING_VIEW)
                TO_STRING_CASE(BINARY)
                TO_STRING_CASE(BINARY_VIEW)
                TO_STRING_CASE(LARGE_STRING)
                TO_STRING_CASE(LARGE_BINARY)
                TO_STRING_CASE(FIXED_SIZE_BINARY)
                TO_STRING_CASE(STRUCT)
                TO_STRING_CASE(LIST)
                TO_STRING_CASE(LARGE_LIST)
                TO_STRING_CASE(LIST_VIEW)
                TO_STRING_CASE(LARGE_LIST_VIEW)
                TO_STRING_CASE(FIXED_SIZE_LIST)
                TO_STRING_CASE(MAP)
                TO_STRING_CASE(DENSE_UNION)
                TO_STRING_CASE(SPARSE_UNION)
                TO_STRING_CASE(DICTIONARY)
                TO_STRING_CASE(RUN_END_ENCODED)
                TO_STRING_CASE(EXTENSION)

#undef TO_STRING_CASE

                default:
                    KLOG(FATAL) << "Unhandled type id: " << id;
                    return "";
            }
        }

        std::string to_string(TimeUnit::type unit) {
            switch (unit) {
                case TimeUnit::SECOND:
                    return "s";
                case TimeUnit::MILLI:
                    return "ms";
                case TimeUnit::MICRO:
                    return "us";
                case TimeUnit::NANO:
                    return "ns";
                default:
                    DKCHECK(false);
                    return "";
            }
        }

    }  // namespace internal

    namespace {

        FieldVector MakeFields(
                std::initializer_list<std::pair<std::string, std::shared_ptr<DataType>>> init_list) {
            FieldVector fields;
            fields.reserve(init_list.size());
            for (const auto &[name, type]: init_list) {
                fields.push_back(field(name, type));
            }
            return fields;
        }

    }  // namespace







    std::ostream &operator<<(std::ostream &os, TimeUnit::type unit) {
        switch (unit) {
            case TimeUnit::SECOND:
                os << "s";
                break;
            case TimeUnit::MILLI:
                os << "ms";
                break;
            case TimeUnit::MICRO:
                os << "us";
                break;
            case TimeUnit::NANO:
                os << "ns";
                break;
        }
        return os;
    }
    // ----------------------------------------------------------------------
    // Schema implementation

    std::string endianness_to_string(Endianness endianness) {
        switch (endianness) {
            case Endianness::Little:
                return "little";
            case Endianness::Big:
                return "big";
            default:
                DKCHECK(false) << "invalid endianness";
                return "???";
        }
    }


    std::shared_ptr<Schema> schema(FieldVector fields,
                                   std::shared_ptr<const KeyValueMetadata> metadata) {
        return std::make_shared<Schema>(std::move(fields), std::move(metadata));
    }

    std::shared_ptr<Schema> schema(
            std::initializer_list<std::pair<std::string, std::shared_ptr<DataType>>> fields,
            std::shared_ptr<const KeyValueMetadata> metadata) {
        return std::make_shared<Schema>(MakeFields(fields), std::move(metadata));
    }

    std::shared_ptr<Schema> schema(FieldVector fields, Endianness endianness,
                                   std::shared_ptr<const KeyValueMetadata> metadata) {
        return std::make_shared<Schema>(std::move(fields), endianness, std::move(metadata));
    }

    std::shared_ptr<Schema> schema(
            std::initializer_list<std::pair<std::string, std::shared_ptr<DataType>>> fields,
            Endianness endianness, std::shared_ptr<const KeyValueMetadata> metadata) {
        return std::make_shared<Schema>(MakeFields(fields), endianness, std::move(metadata));
    }

    turbo::Result<std::shared_ptr<Schema>> unify_schemas(
            const std::vector<std::shared_ptr<Schema>> &schemas,
            const Field::MergeOptions field_merge_options) {
        if (schemas.empty()) {
            return turbo::invalid_argument_error("Must provide at least one schema to unify.");
        }

        if (!schemas[0]->has_distinct_field_names()) {
            return turbo::invalid_argument_error("Can't unify schema with duplicate field names.");
        }

        SchemaBuilder builder{schemas[0], SchemaBuilder::CONFLICT_MERGE, field_merge_options};

        for (size_t i = 1; i < schemas.size(); i++) {
            const auto &schema = schemas[i];
            if (!schema->has_distinct_field_names()) {
                return turbo::invalid_argument_error("Can't unify schema with duplicate field names.");
            }
            TURBO_RETURN_NOT_OK(builder.add_schema(schema));
        }

        return builder.finish();
    }

    const std::shared_ptr<DataType> &utf8_view() {
        static std::shared_ptr<DataType> type = std::make_shared<StringViewType>();
        return type;
    }

    const std::shared_ptr<DataType> &binary_view() {
        static std::shared_ptr<DataType> type = std::make_shared<BinaryViewType>();
        return type;
    }

    std::shared_ptr<DataType> fixed_size_binary(int32_t byte_width) {
        return std::make_shared<FixedSizeBinaryType>(byte_width);
    }

    std::shared_ptr<DataType> duration(TimeUnit::type unit) {
        return std::make_shared<DurationType>(unit);
    }

    std::shared_ptr<DataType> day_time_interval() {
        return std::make_shared<DayTimeIntervalType>();
    }

    std::shared_ptr<DataType> month_day_nano_interval() {
        return std::make_shared<MonthDayNanoIntervalType>();
    }

    std::shared_ptr<DataType> month_interval() {
        return std::make_shared<MonthIntervalType>();
    }

    std::shared_ptr<DataType> timestamp(TimeUnit::type unit) {
        return std::make_shared<TimestampType>(unit);
    }

    std::shared_ptr<DataType> timestamp(TimeUnit::type unit, const std::string &timezone) {
        return std::make_shared<TimestampType>(unit, timezone);
    }

    std::shared_ptr<DataType> time32(TimeUnit::type unit) {
        return std::make_shared<Time32Type>(unit);
    }

    std::shared_ptr<DataType> time64(TimeUnit::type unit) {
        return std::make_shared<Time64Type>(unit);
    }

    std::shared_ptr<DataType> list(std::shared_ptr<DataType> value_type) {
        return std::make_shared<ListType>(std::move(value_type));
    }

    std::shared_ptr<DataType> list(std::shared_ptr<Field> value_field) {
        return std::make_shared<ListType>(std::move(value_field));
    }

    std::shared_ptr<DataType> large_list(std::shared_ptr<DataType> value_type) {
        return std::make_shared<LargeListType>(std::move(value_type));
    }

    std::shared_ptr<DataType> large_list(std::shared_ptr<Field> value_field) {
        return std::make_shared<LargeListType>(std::move(value_field));
    }

    std::shared_ptr<DataType> map(std::shared_ptr<DataType> key_type,
                                  std::shared_ptr<DataType> item_type, bool keys_sorted) {
        return std::make_shared<MapType>(std::move(key_type), std::move(item_type),
                                         keys_sorted);
    }

    std::shared_ptr<DataType> map(std::shared_ptr<DataType> key_type,
                                  std::shared_ptr<Field> item_field, bool keys_sorted) {
        return std::make_shared<MapType>(std::move(key_type), std::move(item_field),
                                         keys_sorted);
    }

    std::shared_ptr<DataType> map(std::shared_ptr<Field> key_field,
                                  std::shared_ptr<Field> item_field, bool keys_sorted) {
        return std::make_shared<MapType>(std::move(key_field), std::move(item_field),
                                         keys_sorted);
    }

    std::shared_ptr<DataType> fixed_size_list(std::shared_ptr<DataType> value_type,
                                              int32_t list_size) {
        return std::make_shared<FixedSizeListType>(std::move(value_type), list_size);
    }

    std::shared_ptr<DataType> fixed_size_list(std::shared_ptr<Field> value_field,
                                              int32_t list_size) {
        return std::make_shared<FixedSizeListType>(std::move(value_field), list_size);
    }

    std::shared_ptr<DataType> list_view(std::shared_ptr<DataType> value_type) {
        return std::make_shared<ListViewType>(std::move(value_type));
    }

    std::shared_ptr<DataType> list_view(std::shared_ptr<Field> value_field) {
        return std::make_shared<ListViewType>(std::move(value_field));
    }

    std::shared_ptr<DataType> large_list_view(std::shared_ptr<DataType> value_type) {
        return std::make_shared<LargeListViewType>(std::move(value_type));
    }

    std::shared_ptr<DataType> large_list_view(std::shared_ptr<Field> value_field) {
        return std::make_shared<LargeListViewType>(std::move(value_field));
    }

    std::shared_ptr<DataType> STRUCT(const FieldVector &fields) {
        return std::make_shared<StructType>(fields);
    }

    std::shared_ptr<DataType> STRUCT(
            std::initializer_list<std::pair<std::string, std::shared_ptr<DataType>>> fields) {
        return std::make_shared<StructType>(MakeFields(fields));
    }

    std::shared_ptr<DataType> run_end_encoded(std::shared_ptr<DataType> run_end_type,
                                              std::shared_ptr<DataType> value_type) {
        return std::make_shared<RunEndEncodedType>(std::move(run_end_type),
                                                   std::move(value_type));
    }

    std::shared_ptr<DataType> sparse_union(FieldVector child_fields,
                                           std::vector<int8_t> type_codes) {
        if (type_codes.empty()) {
            type_codes = internal::Iota(static_cast<int8_t>(child_fields.size()));
        }
        return std::make_shared<SparseUnionType>(std::move(child_fields),
                                                 std::move(type_codes));
    }

    std::shared_ptr<DataType> dense_union(FieldVector child_fields,
                                          std::vector<int8_t> type_codes) {
        if (type_codes.empty()) {
            type_codes = internal::Iota(static_cast<int8_t>(child_fields.size()));
        }
        return std::make_shared<DenseUnionType>(std::move(child_fields), std::move(type_codes));
    }

    FieldVector FieldsFromArraysAndNames(std::vector<std::string> names,
                                         const ArrayVector &arrays) {
        FieldVector fields(arrays.size());
        int i = 0;
        if (names.empty()) {
            for (const auto &array: arrays) {
                fields[i] = field(internal::ToChars(i), array->type());
                ++i;
            }
        } else {
                    DKCHECK_EQ(names.size(), arrays.size());
            for (const auto &array: arrays) {
                fields[i] = field(std::move(names[i]), array->type());
                ++i;
            }
        }
        return fields;
    }

    std::shared_ptr<DataType> sparse_union(const ArrayVector &children,
                                           std::vector<std::string> field_names,
                                           std::vector<int8_t> type_codes) {
        if (type_codes.empty()) {
            type_codes = internal::Iota(static_cast<int8_t>(children.size()));
        }
        auto fields = FieldsFromArraysAndNames(std::move(field_names), children);
        return sparse_union(std::move(fields), std::move(type_codes));
    }

    std::shared_ptr<DataType> dense_union(const ArrayVector &children,
                                          std::vector<std::string> field_names,
                                          std::vector<int8_t> type_codes) {
        if (type_codes.empty()) {
            type_codes = internal::Iota(static_cast<int8_t>(children.size()));
        }
        auto fields = FieldsFromArraysAndNames(std::move(field_names), children);
        return dense_union(std::move(fields), std::move(type_codes));
    }

    std::shared_ptr<DataType> dictionary(const std::shared_ptr<DataType> &index_type,
                                         const std::shared_ptr<DataType> &dict_type,
                                         bool ordered) {
        return std::make_shared<DictionaryType>(index_type, dict_type, ordered);
    }

    std::shared_ptr<Field> field(std::string name, std::shared_ptr<DataType> type,
                                 bool nullable,
                                 std::shared_ptr<const KeyValueMetadata> metadata) {
        return std::make_shared<Field>(std::move(name), std::move(type), nullable,
                                       std::move(metadata));
    }

    std::shared_ptr<Field> field(std::string name, std::shared_ptr<DataType> type,
                                 std::shared_ptr<const KeyValueMetadata> metadata) {
        return std::make_shared<Field>(std::move(name), std::move(type), /*nullable=*/true,
                                       std::move(metadata));
    }

    std::shared_ptr<DataType> decimal(int32_t precision, int32_t scale) {
        return precision <= Decimal128Type::kMaxPrecision ? decimal128(precision, scale)
                                                          : decimal256(precision, scale);
    }

    std::shared_ptr<DataType> decimal128(int32_t precision, int32_t scale) {
        return std::make_shared<Decimal128Type>(precision, scale);
    }

    std::shared_ptr<DataType> decimal256(int32_t precision, int32_t scale) {
        return std::make_shared<Decimal256Type>(precision, scale);
    }

    namespace {

        std::vector<std::shared_ptr<DataType>> g_signed_int_types;
        std::vector<std::shared_ptr<DataType>> g_unsigned_int_types;
        std::vector<std::shared_ptr<DataType>> g_int_types;
        std::vector<std::shared_ptr<DataType>> g_floating_types;
        std::vector<std::shared_ptr<DataType>> g_numeric_types;
        std::vector<std::shared_ptr<DataType>> g_base_binary_types;
        std::vector<std::shared_ptr<DataType>> g_temporal_types;
        std::vector<std::shared_ptr<DataType>> g_interval_types;
        std::vector<std::shared_ptr<DataType>> g_duration_types;
        std::vector<std::shared_ptr<DataType>> g_primitive_types;
        std::once_flag static_data_initialized;

        template<typename T>
        void Extend(const std::vector<T> &values, std::vector<T> *out) {
            out->insert(out->end(), values.begin(), values.end());
        }

        void InitStaticData() {
            // Signed int types
            g_signed_int_types = {int8(), int16(), int32(), int64()};

            // Unsigned int types
            g_unsigned_int_types = {uint8(), uint16(), uint32(), uint64()};

            // All int types
            Extend(g_unsigned_int_types, &g_int_types);
            Extend(g_signed_int_types, &g_int_types);

            // Floating point types
            g_floating_types = {float32(), float64()};

            // Numeric types
            Extend(g_int_types, &g_numeric_types);
            Extend(g_floating_types, &g_numeric_types);

            // Temporal types
            g_temporal_types = {date32(),
                                date64(),
                                time32(TimeUnit::SECOND),
                                time32(TimeUnit::MILLI),
                                time64(TimeUnit::MICRO),
                                time64(TimeUnit::NANO),
                                timestamp(TimeUnit::SECOND),
                                timestamp(TimeUnit::MILLI),
                                timestamp(TimeUnit::MICRO),
                                timestamp(TimeUnit::NANO)};

            // Interval types
            g_interval_types = {day_time_interval(), month_interval(), month_day_nano_interval()};

            // Duration types
            g_duration_types = {duration(TimeUnit::SECOND), duration(TimeUnit::MILLI),
                                duration(TimeUnit::MICRO), duration(TimeUnit::NANO)};

            // Base binary types (without FixedSizeBinary)
            g_base_binary_types = {binary(), utf8(), large_binary(), large_utf8()};

            // Non-parametric, non-nested types. This also DOES NOT include
            //
            // * Decimal
            // * Fixed Size Binary
            // * Time32
            // * Time64
            // * Timestamp
            g_primitive_types = {null(), boolean(), date32(), date64(), binary_view(), utf8_view()};
            Extend(g_numeric_types, &g_primitive_types);
            Extend(g_base_binary_types, &g_primitive_types);
        }

    }  // namespace

    const std::vector<std::shared_ptr<DataType>> &base_binary_types() {
        std::call_once(static_data_initialized, InitStaticData);
        return g_base_binary_types;
    }

    const std::vector<std::shared_ptr<DataType>> &binary_types() {
        static DataTypeVector types = {binary(), large_binary()};
        return types;
    }

    const std::vector<std::shared_ptr<DataType>> &string_types() {
        static DataTypeVector types = {utf8(), large_utf8()};
        return types;
    }

    const std::vector<std::shared_ptr<DataType>> &signed_int_types() {
        std::call_once(static_data_initialized, InitStaticData);
        return g_signed_int_types;
    }

    const std::vector<std::shared_ptr<DataType>> &unsigned_int_types() {
        std::call_once(static_data_initialized, InitStaticData);
        return g_unsigned_int_types;
    }

    const std::vector<std::shared_ptr<DataType>> &int_types() {
        std::call_once(static_data_initialized, InitStaticData);
        return g_int_types;
    }

    const std::vector<std::shared_ptr<DataType>> &floating_point_types() {
        std::call_once(static_data_initialized, InitStaticData);
        return g_floating_types;
    }

    const std::vector<std::shared_ptr<DataType>> &numeric_types() {
        std::call_once(static_data_initialized, InitStaticData);
        return g_numeric_types;
    }

    const std::vector<std::shared_ptr<DataType>> &temporal_types() {
        std::call_once(static_data_initialized, InitStaticData);
        return g_temporal_types;
    }

    const std::vector<std::shared_ptr<DataType>> &interval_types() {
        std::call_once(static_data_initialized, InitStaticData);
        return g_interval_types;
    }

    const std::vector<std::shared_ptr<DataType>> &duration_types() {
        std::call_once(static_data_initialized, InitStaticData);
        return g_duration_types;
    }

    const std::vector<std::shared_ptr<DataType>> &primitive_types() {
        std::call_once(static_data_initialized, InitStaticData);
        return g_primitive_types;
    }

    const std::vector<TimeUnit::type> &TimeUnit::values() {
        static std::vector<TimeUnit::type> units = {TimeUnit::SECOND, TimeUnit::MILLI,
                                                    TimeUnit::MICRO, TimeUnit::NANO};
        return units;
    }

}  // namespace nebula
